#!/usr/bin/env python3
"""
网络分析仪控制器
重构版本 - 面向对象设计
"""

import pyvisa
import time
import csv
from datetime import datetime
from typing import Optional, Dict, Any, List
import logging


class NetworkAnalyzerError(Exception):
    """网络分析仪相关错误的基类"""
    pass


class ConnectionError(NetworkAnalyzerError):
    """连接错误"""
    pass


class ConfigurationError(NetworkAnalyzerError):
    """配置错误"""
    pass


class MeasurementError(NetworkAnalyzerError):
    """测量错误"""
    pass


class InstrumentWrapper:
    """
    仪器连接包装器
    提供自动重连机制和连接状态监控
    """
    
    def __init__(self, resource_manager, resource_string: str, timeout: int, logger):
        """
        初始化仪器包装器
        
        Args:
            resource_manager: VISA资源管理器
            resource_string: 资源连接字符串
            timeout: 通信超时时间(毫秒)
            logger: 日志记录器
        """
        self.resource_manager = resource_manager
        self.resource_string = resource_string
        self.timeout = timeout
        self.logger = logger
        self._instrument = None
        self._max_reconnect_attempts = 3
        self._reconnect_delay = 2.0  # 重连延迟(秒)
        self._last_command_time = 0
        self._command_timeout = 30.0  # 命令超时(秒)
        
    @property
    def is_connected(self) -> bool:
        """检查连接状态"""
        if self._instrument is None:
            return False
        try:
            # 通过查询设备ID来验证连接
            self._instrument.query("*IDN?", delay=0.1)
            return True
        except Exception:
            return False
    
    def _reconnect(self) -> bool:
        """
        执行重连操作
        
        Returns:
            bool: 重连成功返回True
        """
        self.logger.warning("检测到连接中断，尝试重连...")
        
        # 关闭现有连接
        if self._instrument is not None:
            try:
                self._instrument.close()
            except Exception as e:
                self.logger.debug(f"关闭旧连接时出错: {e}")
        
        # 尝试重连
        for attempt in range(self._max_reconnect_attempts):
            try:
                self.logger.info(f"重连尝试 {attempt + 1}/{self._max_reconnect_attempts}")
                
                # 重新创建资源管理器（有时需要重新初始化）
                try:
                    self.resource_manager.close()
                except:
                    pass
                
                import pyvisa
                self.resource_manager = pyvisa.ResourceManager()
                self._instrument = self.resource_manager.open_resource(self.resource_string)
                self._instrument.timeout = self.timeout
                
                # 验证连接 - 使用超时保护
                self._instrument.timeout = 5000  # 验证时使用较短超时
                idn = self._instrument.query("*IDN?").strip()
                self._instrument.timeout = self.timeout  # 恢复原始超时
                
                self.logger.info(f"重连成功: {idn}")
                return True
                
            except Exception as e:
                self.logger.warning(f"重连尝试 {attempt + 1} 失败: {e}")
                if attempt < self._max_reconnect_attempts - 1:
                    time.sleep(self._reconnect_delay * (attempt + 1))
                
                # 在最后一次尝试前，尝试更激进的重置
                if attempt == self._max_reconnect_attempts - 2:
                    self.logger.warning("尝试重新初始化VISA资源...")
                    try:
                        import pyvisa
                        pyvisa.ResourceManager().close()
                    except:
                        pass
        
        self.logger.error("重连失败，已达到最大重试次数")
        return False
    
    def _execute_with_retry(self, operation, *args, **kwargs):
        """
        执行操作并支持自动重试
        
        Args:
            operation: 要执行的操作函数
            *args, **kwargs: 操作参数
            
        Returns:
            操作返回结果
            
        Raises:
            最后一次操作异常
        """
        last_exception = None
        
        for attempt in range(self._max_reconnect_attempts + 1):
            try:
                if not self.is_connected:
                    if not self._reconnect():
                        raise ConnectionError("无法建立连接")
                
                self._last_command_time = time.time()
                return operation(*args, **kwargs)
                
            except (pyvisa.VisaIOError, ConnectionError) as e:
                last_exception = e
                self.logger.warning(f"操作失败，尝试 {attempt + 1}/{self._max_reconnect_attempts + 1}: {e}")
                
                # 处理各种VISA错误类型
                if isinstance(e, pyvisa.VisaIOError):
                    error_code = getattr(e, 'error_code', None)
                    if error_code in [pyvisa.constants.VI_ERROR_CONN_LOST, 
                                     pyvisa.constants.VI_ERROR_IO,
                                     pyvisa.constants.VI_ERROR_TMO]:
                        # 对于I/O错误或超时，强制重新连接
                        self.logger.warning("检测到严重通信错误，强制重新连接...")
                        if self._instrument is not None:
                            try:
                                self._instrument.close()
                            except:
                                pass
                            self._instrument = None
                        
                        if not self._reconnect():
                            break
                
                elif attempt < self._max_reconnect_attempts:
                    time.sleep(self._reconnect_delay * (attempt + 1))  # 指数退避
        
        raise last_exception
    
    def write(self, command: str) -> None:
        """发送写命令，支持自动重连"""
        return self._execute_with_retry(
            lambda cmd: self._instrument.write(cmd),
            command
        )
    
    def read(self) -> str:
        """读取数据，支持自动重连"""
        return self._execute_with_retry(
            lambda: self._instrument.read()
        )
    
    def query(self, command: str, delay: float = None) -> str:
        """发送查询命令，支持自动重连"""
        return self._execute_with_retry(
            lambda cmd, delay=None: self._instrument.query(cmd, delay=delay),
            command,
            delay=delay
        )
    
    def query_ascii_values(self, command: str, **kwargs):
        """发送查询命令并获取ASCII值，支持自动重连"""
        return self._execute_with_retry(
            lambda cmd, **kw: self._instrument.query_ascii_values(cmd, **kw),
            command,
            **kwargs
        )
    
    def close(self) -> None:
        """关闭连接"""
        if self._instrument is not None:
            try:
                self._instrument.close()
                self.logger.info("仪器连接已关闭")
            except Exception as e:
                self.logger.warning(f"关闭连接时出错: {e}")
            finally:
                self._instrument = None


class NetworkAnalyzerController:
    """
    网络分析仪控制器类
    提供对Ceyear 3671G网络分析仪的完整控制功能，包含自动重连机制
    """
    
    def __init__(self, ip_address: str = "192.168.1.10", timeout: int = 10000):
        """
        初始化网络分析仪控制器
        
        Args:
            ip_address: 网络分析仪IP地址
            timeout: 通信超时时间(毫秒)
        """
        self.ip_address = ip_address
        self.timeout = timeout
        self.resource_manager = None
        self.instrument = None  # 现在是InstrumentWrapper实例
        self.logger = self._setup_logger()
        self._connection_attempts = 0
        self._max_initial_attempts = 3
        
    def _setup_logger(self) -> logging.Logger:
        """设置日志记录器"""
        logger = logging.getLogger(__name__)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger
    
    def connect(self) -> bool:
        """
        连接到网络分析仪，支持自动重连
        
        Returns:
            bool: 连接成功返回True
            
        Raises:
            ConnectionError: 连接失败时抛出
        """
        try:
            self.resource_manager = pyvisa.ResourceManager()
            resources = self.resource_manager.list_resources()
            self.logger.info(f"可用资源: {resources}")
            
            self.resource_string = f"TCPIP0::{self.ip_address}::inst0::INSTR"
            self.instrument = InstrumentWrapper(
                self.resource_manager, 
                self.resource_string, 
                self.timeout, 
                self.logger
            )
            
            # 验证连接
            idn = self.instrument.query('*IDN?').strip()
            self.logger.info(f"设备连接成功: {idn}")
            return True
            
        except Exception as e:
            self.logger.error(f"连接失败: {e}")
            raise ConnectionError(f"无法连接到网络分析仪: {e}")
    
    def disconnect(self) -> None:
        """断开与网络分析仪的连接"""
        if self.instrument:
            try:
                self.instrument.close()
                self.logger.info("已断开与网络分析仪的连接")
            except Exception as e:
                self.logger.warning(f"断开连接时出错: {e}")
        
        if self.resource_manager:
            try:
                self.resource_manager.close()
            except Exception as e:
                self.logger.warning(f"关闭资源管理器时出错: {e}")
        
        self.logger.info("已断开连接")
    
    def __enter__(self):
        """上下文管理器入口"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.disconnect()
    
    def reset(self) -> bool:
        """
        重置仪器到默认状态
        
        Returns:
            bool: 重置成功返回True
        """
        try:
            self.logger.info("重置仪器...")
            self.instrument.write('*RST')
            time.sleep(5)  # 等待重置完成
            self.instrument.write('*CLS')  # 清除状态寄存器
            time.sleep(1)
            self.logger.info("仪器重置完成")
            return True
        except Exception as e:
            self.logger.error(f"重置失败: {e}")
            return False
    
    def check_error(self) -> Optional[str]:
        """
        检查系统错误
        
        Returns:
            str: 错误信息，无错误返回None
        """
        try:
            error = self.instrument.query('SYSTem:ERRor?').strip()
            if not error.startswith('+0'):
                self.logger.error(f"系统错误: {error}")
                return error
            return None
        except Exception as e:
            self.logger.error(f"查询错误失败: {e}")
            return str(e)
    
    def configure_frequency_range(self, start_freq: float, stop_freq: float, 
                                  points: int = 201) -> bool:
        """
        配置频率范围
        
        Args:
            start_freq: 起始频率(Hz)
            stop_freq: 停止频率(Hz)
            points: 扫描点数
            
        Returns:
            bool: 配置成功返回True
            
        Raises:
            ConfigurationError: 配置失败时抛出
        """
        try:
            self.logger.info(f"配置频率范围: {start_freq/1e9:.2f}-{stop_freq/1e9:.2f}GHz, {points}点")
            
            # 设置扫描类型
            self.instrument.write('SENSe1:SWEep:TYPE LINear')
            time.sleep(0.5)
            
            # 设置频率范围
            self.instrument.write(f'SENSe1:FREQuency:STARt {start_freq}')
            time.sleep(0.5)
            self.instrument.write(f'SENSe1:FREQuency:STOP {stop_freq}')
            time.sleep(0.5)
            self.instrument.write(f'SENSe1:SWEep:POINts {points}')
            time.sleep(0.5)
            
            # 验证设置
            actual_start = float(self.instrument.query('SENSe1:FREQuency:STARt?'))
            actual_stop = float(self.instrument.query('SENSe1:FREQuency:STOP?'))
            actual_points = int(self.instrument.query('SENSe1:SWEep:POINts?'))
            
            self.logger.info(f"实际设置: {actual_start/1e9:.2f}-{actual_stop/1e9:.2f}GHz, {actual_points}点")
            
            # 设置连续扫描模式
            self.instrument.write('SENSe1:SWEep:MODE CONTinuous')
            time.sleep(0.5)
            
            return True
            
        except Exception as e:
            self.logger.error(f"频率范围配置失败: {e}")
            raise ConfigurationError(f"频率范围配置失败: {e}")
    
    def configure_display(self, window_num: int = 1, split_mode: str = "D12") -> bool:
        """
        配置显示窗口
        
        Args:
            window_num: 窗口编号
            split_mode: 分割模式 (D12: 上下分割)
            
        Returns:
            bool: 配置成功返回True
        """
        try:
            self.logger.info(f"配置显示窗口{window_num}")
            
            # 确保窗口开启
            self.instrument.write(f'DISPlay:WINDow{window_num}:STATe ON')
            time.sleep(0.5)
            
            # 设置窗口布局
            self.instrument.write(f'DISPlay:WINDow{window_num}:SPLit {split_mode}')
            time.sleep(0.5)
            
            self.logger.info("显示窗口配置完成")
            return True
            
        except Exception as e:
            self.logger.error(f"显示窗口配置失败: {e}")
            return False
    
    def create_trace(self, trace_name: str, parameter: str, 
                     channel: int = 1, trace_type: str = "MLOGarithmic") -> bool:
        """
        创建测量轨迹
        
        Args:
            trace_name: 轨迹名称
            parameter: 测量参数 (如S21)
            channel: 通道号
            trace_type: 轨迹类型 (MLOGarithmic, PHASe等)
            
        Returns:
            bool: 创建成功返回True
        """
        try:
            self.logger.info(f"创建轨迹: {trace_name} ({parameter}) 通道{channel}")
            
            # 删除所有现有参数
            self.instrument.write('CALCulate:PARameter:DELete:ALL')
            time.sleep(1)
            
            # 创建参数
            self.instrument.write(f'CALCulate{channel}:PARameter:DEFine "{trace_name}",{parameter}')
            time.sleep(1)
            
            # 选择参数
            self.instrument.write(f'CALCulate{channel}:PARameter:SELect "{trace_name}"')
            time.sleep(0.5)
            
            # 设置格式
            self.instrument.write(f'CALCulate{channel}:FORMat {trace_type}')
            time.sleep(0.5)
            
            # 分配到显示窗口
            trace_num = 1 if channel == 1 else 2
            self.instrument.write(f'DISPlay:WINDow1:TRACe{trace_num}:FEED "{trace_name}"')
            time.sleep(0.5)
            
            self.instrument.write(f'DISPlay:WINDow1:TRACe{trace_num}:STATe ON')
            time.sleep(0.5)
            
            self.logger.info(f"轨迹 {trace_name} 创建完成")
            return True
            
        except Exception as e:
            self.logger.error(f"轨迹创建失败: {e}")
            return False
    
    def get_marker_data(self, channel: int, frequency: float) -> Dict[str, float]:
        """
        获取指定频率的标记数据
        
        Args:
            channel: 通道号
            frequency: 目标频率(Hz)
            
        Returns:
            dict: 包含频率、幅度、相位的数据字典
            
        Raises:
            MeasurementError: 测量失败时抛出
        """
        try:
            # 选择参数
            self.instrument.write(f'CALCulate{channel}:PARameter:SELect "S21_MAG"')
            time.sleep(0.2)
            
            # 开启标记
            self.instrument.write(f'CALCulate{channel}:MARKer1:STATe ON')
            time.sleep(0.2)
            
            # 设置标记频率
            self.instrument.write(f'CALCulate{channel}:MARKer1:X {frequency}')
            time.sleep(0.2)
            
            # 读取频率
            actual_freq = float(self.instrument.query(f'CALCulate{channel}:MARKer1:X?'))
            
            # 读取幅度
            magnitude_str = self.instrument.query(f'CALCulate{channel}:MARKer1:Y?').strip()
            if ',' in magnitude_str:
                magnitude = float(magnitude_str.split(',')[0])
            else:
                magnitude = float(magnitude_str)
            
            # 读取相位
            self.instrument.write(f'CALCulate{channel}:FORMat PHASe')
            time.sleep(0.2)
            phase_str = self.instrument.query(f'CALCulate{channel}:MARKer1:Y?').strip()
            if ',' in phase_str:
                phase = float(phase_str.split(',')[0])
            else:
                phase = float(phase_str)
            
            return {
                'frequency': actual_freq,
                'magnitude': magnitude,
                'phase': phase
            }
            
        except Exception as e:
            self.logger.error(f"标记数据获取失败: {e}")
            raise MeasurementError(f"标记数据获取失败: {e}")
    
    def auto_scale_trace(self, trace_num: int = 1) -> bool:
        """
        自动缩放轨迹
        
        Args:
            trace_num: 轨迹编号
            
        Returns:
            bool: 操作成功返回True
        """
        try:
            self.instrument.write(f'DISPlay:WINDow1:TRACe{trace_num}:Y:AUTO')
            time.sleep(0.5)
            return True
        except Exception as e:
            self.logger.error(f"自动缩放失败: {e}")
            return False


class MeasurementSession:
    """
    测量会话管理类
    负责管理一次完整的测量会话
    """
    
    def __init__(self, controller: NetworkAnalyzerController, 
                 target_frequency: float = 5.44e9):
        """
        初始化测量会话
        
        Args:
            controller: 网络分析仪控制器实例
            target_frequency: 目标测量频率(Hz)
        """
        self.controller = controller
        self.target_frequency = target_frequency
        self.measurement_count = 0
        self.csv_filename = None
        self.csv_file = None
        self.csv_writer = None
        
    def start_session(self, output_dir: str = ".") -> bool:
        """
        开始测量会话
        
        Args:
            output_dir: 输出文件目录
            
        Returns:
            bool: 会话启动成功返回True
        """
        try:
            # 创建CSV文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.csv_filename = f"{output_dir}/measurement_results_{timestamp}.csv"
            
            # 打开CSV文件
            self.csv_file = open(self.csv_filename, 'w', newline='', encoding='utf-8')
            fieldnames = ['测试序号', '测试时间', '频率(GHz)', '幅度(dB)', 
                         '相位(度)', '信号质量', '目标频率(GHz)']
            self.csv_writer = csv.DictWriter(self.csv_file, fieldnames=fieldnames)
            self.csv_writer.writeheader()
            
            # 打印表头
            print("\n" + "=" * 100)
            print(f"{'测试序号':<8} {'测试时间':<20} {'频率(GHz)':<12} "
                  f"{'幅度(dB)':<12} {'相位(度)':<12} {'信号质量':<15}")
            print("=" * 100)
            
            return True
            
        except Exception as e:
            self.controller.logger.error(f"会话启动失败: {e}")
            return False
    
    def stop_session(self):
        """停止测量会话"""
        if self.csv_file:
            self.csv_file.close()
            self.controller.logger.info(f"测量数据已保存到: {self.csv_filename}")
    
    def perform_measurement(self) -> Dict[str, Any]:
        """
        执行一次测量
        
        Returns:
            dict: 测量结果数据
        """
        try:
            self.measurement_count += 1
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            
            # 获取测量数据
            data = self.controller.get_marker_data(1, self.target_frequency)
            
            # 判断信号质量
            magnitude = abs(data['magnitude'])
            if magnitude < 3:
                signal_quality = "良好"
            elif magnitude < 10:
                signal_quality = "中等"
            else:
                signal_quality = "较差"
            
            # 格式化数据
            result = {
                '测试序号': self.measurement_count,
                '测试时间': current_time,
                '频率(GHz)': data['frequency'] / 1e9,
                '幅度(dB)': data['magnitude'],
                '相位(度)': data['phase'],
                '信号质量': signal_quality,
                '目标频率(GHz)': self.target_frequency / 1e9
            }
            
            # 输出到控制台
            print(f"{result['测试序号']:<8} {result['测试时间']:<20} "
                  f"{result['频率(GHz)']:<12.6f} {result['幅度(dB)']:<12.3f} "
                  f"{result['相位(度)']:<12.1f} {result['信号质量']:<15}")
            
            # 保存到CSV
            if self.csv_writer:
                self.csv_writer.writerow(result)
            
            return result
            
        except Exception as e:
            self.controller.logger.error(f"测量失败: {e}")
            return {
                '测试序号': self.measurement_count,
                '测试时间': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                '频率(GHz)': self.target_frequency / 1e9,
                '幅度(dB)': 'N/A',
                '相位(度)': 'N/A',
                '信号质量': '错误',
                '目标频率(GHz)': self.target_frequency / 1e9
            }
    
    def run_continuous_measurement(self, interval: float = 3.0, reconnect_before_start: bool = False):
        """
        运行连续测量
        
        Args:
            interval: 测量间隔(秒)
            reconnect_before_start: 开始前是否断开重连
        """
        try:
            # 如果要求重连，先断开再连接
            if reconnect_before_start:
                self.controller.logger.info("执行断开重连操作...")
                self.controller.disconnect()
                time.sleep(2)  # 等待完全断开
                
                if not self.controller.connect():
                    self.controller.logger.error("重连失败，无法开始连续测量")
                    return
                
                self.controller.logger.info("重连成功，继续连续测量")
            
            self.controller.logger.info("开始连续测量模式")
            print("按 Ctrl+C 停止测试")
            
            while True:
                self.perform_measurement()
                time.sleep(interval)
                
        except KeyboardInterrupt:
            self.controller.logger.info("用户停止测量")
            print(f"\n\n用户停止测试 (共进行了 {self.measurement_count} 次测试)")
            print("=" * 100)


def main():
    """
    主函数 - 演示如何使用重构后的类
    """
    try:
        # 使用上下文管理器确保正确关闭连接
        with NetworkAnalyzerController("192.168.1.10") as nwa:
            # 重置仪器
            nwa.reset()
            
            # 配置频率范围
            nwa.configure_frequency_range(5.2e9, 5.8e9, 201)
            
            # 配置显示
            nwa.configure_display(1, "D12")
            
            # 创建轨迹
            nwa.create_trace("S21_MAG", "S21", 1, "MLOGarithmic")
            nwa.create_trace("S21_PHASE", "S21", 2, "PHASe")
            
            # 自动缩放轨迹
            nwa.auto_scale_trace(1)
            nwa.auto_scale_trace(2)
            
            # 开始测量会话
            session = MeasurementSession(nwa, 5.44e9)
            if session.start_session():
                try:
                    # 运行连续测量
                    session.run_continuous_measurement(3.0)
                finally:
                    session.stop_session()
                    
    except NetworkAnalyzerError as e:
        print(f"网络分析仪错误: {e}")
    except Exception as e:
        print(f"意外错误: {e}")


if __name__ == "__main__":
    main()